from Crypto.Cipher import AES

class AESCipher:
    def __init__(self, key: bytes, iv: bytes, padding: str = "none"):
        self.key = key
        self.iv = iv
        self.padding = padding.lower()
    
    def _pad(self, data: bytes) -> bytes:
        if self.padding == "pkcs7":
            """PKCS7填充"""
            pad_len = 16 - len(data) % 16
            return data + bytes([pad_len]) * pad_len
        elif self.padding == "zero":
            """零填充"""
            pad_len = 16 - len(data) % 16
            return data + bytes([0]) * pad_len
        elif self.padding == "none":
            return data
        else:
            raise ValueError(f"Unsupported padding: {self.padding}")
        
    def _unpad(self, data: bytes) -> bytes:
        if self.padding == "pkcs7": 
            """PKCS7去填充"""
            pad_len = data[-1]
            return data[:-pad_len]
        elif self.padding == "zero":
            """零填充去填充"""
            return data.rstrip(b'\x00')
        elif self.padding == "none":
            return data
        else:
            raise ValueError(f"Unsupported padding: {self.padding}")
    
    def encrypt(self, data: str) -> str:
        """AES加密"""
        padded_data = self._pad(data.encode('utf-8'))
        cipher = AES.new(self.key, AES.MODE_CBC, self.iv)
        encrypted_bytes = cipher.encrypt(padded_data)
        """返回16进制字符串"""
        return encrypted_bytes.hex()
    
    def decrypt(self, data: str) -> str:
        """AES解密"""
        """将16进制字符串转换为字节"""
        encrypted_bytes = bytes.fromhex(data)
        cipher = AES.new(self.key, AES.MODE_CBC, self.iv)
        padded_data = cipher.decrypt(encrypted_bytes)
        return self._unpad(padded_data).decode('utf-8')
